package com.kklsqm.langchain4j.service;

import cn.hutool.json.JSON;
import cn.hutool.json.JSONUtil;
import okhttp3.*;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.Map;

@Service
public class ImageInstanceSegmentationService {

    private static final String API_URL = "https://dashscope.aliyuncs.com/api/v1/services/aigc/image2image/image-synthesis";
    private static final String QUERY_URL_TEMPLATE = "https://dashscope.aliyuncs.com/api/v1/tasks/%s";
    private static final String API_KEY = "sk-072acf6e9113491bb5809110bd83e8b1"; // 替换成你的真实Key

    private final OkHttpClient client = new OkHttpClient();

    public Map<String, String> segmentImage(String imageUrl) {
        try {
            // 构造请求体
            JSON input = JSONUtil.createObj()
                    .put("image_url", imageUrl);

            JSON parameters = JSONUtil.createObj(); // 根据需要可以添加额外的参数

            JSON body = JSONUtil.createObj()
                    .put("model", "image-instance-segmentation")
                    .put("input", input)
                    .put("parameters", parameters);

            // Step 1: 创建任务
            Request request = new Request.Builder()
                    .url(API_URL)
                    .addHeader("Authorization", "Bearer " + API_KEY)
                    .addHeader("Content-Type", "application/json")
                    .addHeader("X-DashScope-Async", "enable")
                    .post(RequestBody.create(body.toString(), MediaType.parse("application/json")))
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    String errorBody = response.body().string();
                    throw new IOException("创建任务失败，HTTP状态码: " + response.code() + "\n错误信息: " + errorBody);
                }

                JSON json = JSONUtil.parseObj(response.body().string());
                String taskId = json.getByPath("output.task_id", String.class);

                // Step 2: 轮询任务结果
                return pollTaskResult(taskId);
            }
        } catch (Exception e) {
            throw new RuntimeException("调用人物实例分割服务失败: " + e.getMessage(), e);
        }
    }

    private Map<String, String> pollTaskResult(String taskId) throws IOException, InterruptedException {
        String queryUrl = String.format(QUERY_URL_TEMPLATE, taskId);
        Request queryRequest = new Request.Builder()
                .url(queryUrl)
                .addHeader("Authorization", "Bearer " + API_KEY)
                .get()
                .build();

        while (true) {
            try (Response queryResponse = client.newCall(queryRequest).execute()) {
                if (!queryResponse.isSuccessful()) {
                    throw new IOException("查询任务状态失败，HTTP状态码: " + queryResponse.code());
                }

                JSON result = JSONUtil.parseObj(queryResponse.body().string());
                JSON output = result.getByPath("output", JSON.class);
                String status = (String) output.getByPath("task_status");

                System.out.println("当前任务状态: " + status);

                if ("SUCCEEDED".equals(status)) {
                    // output_vis_image_url

                    return Map.of("output_image_url", (String) output.getByPath("output_image_url"), "output_vis_image_url", (String) output.getByPath("output_vis_image_url"));
                } else if ("FAILED".equals(status)) {
                    throw new RuntimeException("人物实例分割任务执行失败: " + result.toStringPretty());
                }

                Thread.sleep(3000); // 等待3秒后重试
            }
        }
    }
}
